import csv
from datetime import datetime
from typing import List, Tuple, Optional

from pytz import timezone

from vnpy.trader.engine import BaseEngine, MainEngine, EventEngine
from vnpy.trader.constant import Interval, Exchange
from vnpy.trader.object import BarData, HistoryRequest, FutBasicRequest
from vnpy.feeds.database import BaseDatabase, get_database, BarOverview
from vnpy.feeds.datafeed import BaseDatafeed, get_datafeed

import pandas as pd
from pandas import DataFrame

# APP_NAME = "DataManager"

EXCHANGE_LIST = [
    "SH",
    "SZ"
]

EXCHANGE_TS2VT = {
    "CFX": Exchange.CFFEX,
    "SHF": Exchange.SHFE,
    "ZCE": Exchange.CZCE,
    "DCE": Exchange.DCE,
    "INE": Exchange.INE,
    "SH": Exchange.SSE,
    "SZ": Exchange.SZSE,
}


def to_vt_symbol(symbol, exchange) -> Optional[str]:
    """将tushare转换为交易所代码"""
    if exchange in EXCHANGE_LIST:
        ts_symbol = f"{symbol}.{EXCHANGE_TS2VT[exchange].value}".upper()
    else:
        return None
    return ts_symbol


def to_vt_exchange(exchange) -> Optional[str]:
    """将tushare转换为交易所代码"""
    if exchange.upper() in EXCHANGE_LIST:
        vt_exchange = EXCHANGE_TS2VT[exchange]
    else:
        return None

    return vt_exchange


class DataManager:  # BaseEngine
    """"""

    def __init__(self, database_type: str = "mongodb", database_name: str = "test",
                 datafeed_name: str = "tushare"):

        # main_engine: MainEngine,
        # event_engine: EventEngine,
        """"""
        # super().__init__(main_engine, event_engine, APP_NAME)

        self.database: BaseDatabase = get_database(database_type, database_name)
        self.datafeed: BaseDatafeed = get_datafeed(datafeed_name)

    def import_data_from_csv(
            self,
            file_path: str,
            symbol_head: str,
            exchange: Exchange,
            interval: Interval,
            tz_name: str,
            datetime_head: str,
            open_head: str,
            high_head: str,
            low_head: str,
            close_head: str,
            volume_head: str,
            turnover_head: str,
            open_interest_head: str,
            datetime_format: str
    ) -> Tuple:
        """"""
        with open(file_path, "rt", encoding='utf-8-sig') as f:
            buf = [line.replace("\0", "") for line in f]

        reader = csv.DictReader(buf, delimiter=",")

        bars = []
        start = None
        count = 0
        tz = timezone(tz_name)

        for item in reader:
            if datetime_format:
                dt = datetime.strptime(item[datetime_head], datetime_format)
            else:
                dt = datetime.fromisoformat(item[datetime_head])

            dt = tz.localize(dt)

            ts_code: str = item.get(symbol_head, 0)
            symbol = ts_code.split(".")[0]
            exchange = ts_code.split(".")[-1]

            exchange = to_vt_exchange(exchange)

            turnover = item.get(turnover_head, 0)

            # if open_interest_head:
            # open_interest = item.get(open_interest_head, 0)
            open_interest = 0
            bar = BarData(
                symbol=symbol,
                exchange=exchange,
                datetime=dt,
                interval=interval,
                volume=float(item[volume_head]),
                open_price=float(item[open_head]),
                high_price=float(item[high_head]),
                low_price=float(item[low_head]),
                close_price=float(item[close_head]),
                turnover=float(turnover),
                open_interest=float(open_interest),
                gateway_name="DB",
            )

            bars.append(bar)

            # do some statistics
            count += 1
            if not start:
                start = bar.datetime

        end = bar.datetime

        # insert into database
        self.database.save_bar_data(bars)

        return start, end, count

    def output_data_to_csv(
            self,
            file_path: str,
            symbol: str,
            exchange: Exchange,
            interval: Interval,
            start: datetime,
            end: datetime
    ) -> bool:
        """"""
        bars = self.load_bar_data(symbol, exchange, interval, start, end)

        fieldnames = [
            "symbol",
            "exchange",
            "datetime",
            "open",
            "high",
            "low",
            "close",
            "volume",
            "turnover",
            "open_interest"
        ]

        try:
            with open(file_path, "w") as f:
                writer = csv.DictWriter(f, fieldnames=fieldnames, lineterminator="\n")
                writer.writeheader()

                for bar in bars:
                    d = {
                        "symbol": bar.symbol,
                        "exchange": bar.exchange.value,
                        "datetime": bar.datetime.strftime("%Y-%m-%d %H:%M:%S"),
                        "open": bar.open_price,
                        "high": bar.high_price,
                        "low": bar.low_price,
                        "close": bar.close_price,
                        "turnover": bar.turnover,
                        "volume": bar.volume,
                        "open_interest": bar.open_interest,
                    }
                    writer.writerow(d)

            return True
        except PermissionError:
            return False

    def get_bar_overview(self) -> List[BarOverview]:
        """"""
        return self.database.get_bar_overview()

    def load_bar_data(
            self,
            symbol: str,
            exchange: Exchange,
            interval: Interval,
            start: datetime,
            end: datetime
    ) -> List[BarData]:
        """"""
        bars = self.database.load_bar_data(
            symbol,
            exchange,
            interval,
            start,
            end
        )

        return bars

    def delete_bar_data(
            self,
            symbol: str,
            exchange: Exchange,
            interval: Interval
    ) -> int:
        """"""
        count = self.database.delete_bar_data(
            symbol,
            exchange,
            interval
        )

        return count

    def download_bar_data(
            self,
            symbol: str,
            exchange: Exchange,
            interval: str,
            start: datetime,
            end: datetime = datetime.now()  # DB_TZ
    ) -> int:
        """
        Query bar data from datafeed.
        """
        req = HistoryRequest(
            symbol=symbol,
            exchange=exchange,
            interval=Interval(interval),
            start=start,
            end=end
        )

        vt_symbol = f"{symbol}.{exchange.value}"
        # contract = self.main_engine.get_contract(vt_symbol)
        #
        # # If history data provided in gateway, then query
        # if contract and contract.history_data:
        #     data = self.main_engine.query_history(
        #         req, contract.gateway_name
        #     )
        # # Otherwise use datafeed to query data
        # else:
        data = self.datafeed.query_bar_history(req)

        if len(data):
            self.database.save_bar_data(data)
            return (len(data))

        return 0

    def download_tick_data(
            self,
            symbol: str,
            exchange: Exchange,
            start: datetime
    ) -> int:
        """
        Query tick data from datafeed.
        """
        req = HistoryRequest(
            symbol=symbol,
            exchange=exchange,
            start=start,
            end=datetime.now(DB_TZ)
        )

        data = self.datafeed.query_tick_history(req)

        if data:
            self.database.save_tick_data(data)
            return (len(data))

        return 0

    def download_fut_basic(self, exchange: Exchange = "DCE") -> int:
        """
        Query fut basic from datafeed.
        """
        req = FutBasicRequest(
            exchange=exchange,
            fut_type=1
        )
        data = self.datafeed.query_fut_basic(req)

        if data:
            self.database.save_fut_basic(data)
            return (len(data))

        return 0

    def download_fut_mapping(self, ts_code: str = "SAL.ZCE", start_date: str = "start_date",
                             end_date: str = "end_date") -> int:
        """
        Query fut basic from datafeed.
        """
        # req = FutBasicRequest(
        #     exchange=exchange,
        #     fut_type=1
        # )
        data = self.datafeed.query_fut_mapping(ts_code, start_date, end_date)

        if data:
            self.database.save_fut_mapping(data)
            return (len(data))

        return 0

    def get_all_securities(self) -> DataFrame:
        data = self.database.get_all_securities()

        if len(data):
            return data
        else:
            data = self.datafeed.get_all_securities()
            self.database.save_all_securities(data)
            return data